# 48/100 二叉树-路径总合III
# leetcode第437题: https://leetcode.cn/problems/path-sum-iii/description/?envType=study-plan-v2&envId=top-100-liked
# Date: 2025/1/8
from typing import Optional

from leetcode.bds import TreeNode, TreeConverter
import leetcode.test as test


def pathSum(root: Optional[TreeNode], targetSum: int) -> int:
    """我的方法, 使用中序遍历来遍历所有节点, 在中序遍历中使用回溯法来检查所有可能性
    缺点是计算复杂度过高"""

    if not root:
        return 0

    def backtracking(path: list[TreeNode], sum_num: int):
        """回溯法, 检查以path[0]为起始节点的所有可能性"""
        nonlocal res
        if sum_num == targetSum:
            res += 1

        if path[-1].left:
            sum_num += path[-1].left.val
            path.append(path[-1].left)
            backtracking(path, sum_num)
            pop = path.pop()
            sum_num -= pop.val
        if path[-1].right:
            sum_num += path[-1].right.val
            path.append(path[-1].right)
            backtracking(path, sum_num)
            pop = path.pop()
            sum_num -= pop.val

    def in_order_travel(curr: TreeNode):
        """中序遍历"""
        if curr.left:
            in_order_travel(curr.left)
        backtracking([curr], curr.val)
        if curr.right:
            in_order_travel(curr.right)

    res = 0
    in_order_travel(root)
    return res


def pathSum_opt(root: Optional[TreeNode], targetSum: int) -> int:
    def count_paths(node: TreeNode, current_sum: int, prefix_sums: dict) -> int:
        """
        使用前缀和计算从根节点到当前节点的路径和。
        """
        if not node:
            return 0

        # 当前路径和
        current_sum += node.val

        # 从前缀和中找到当前路径和 - targetSum 的次数
        count = prefix_sums.get(current_sum - targetSum, 0)

        # 更新前缀和字典
        prefix_sums[current_sum] = prefix_sums.get(current_sum, 0) + 1

        # 递归计算左右子树
        count += count_paths(node.left, current_sum, prefix_sums)
        count += count_paths(node.right, current_sum, prefix_sums)

        # 回溯时移除当前路径和
        prefix_sums[current_sum] -= 1
        if prefix_sums[current_sum] == 0:
            del prefix_sums[current_sum]

        return count

    # 初始状态的前缀和字典
    prefix_sums = {0: 1}
    return count_paths(root, 0, prefix_sums)


if __name__ == '__main__':
    root1 = TreeConverter.list_to_tree([10, 5, -3, 3, 2, None, 11, 3, -2, None, 1])
    root2 = TreeConverter.list_to_tree([5, 4, 8, 11, None, 13, 4, 7, 2, None, None, 5, 1])
    # TreeConverter.draw_tree_from_root(root2)

    inp = [{"root": root1, "targetSum": 8},
           {"root": root2, "targetSum": 22}, ]
    out = [3, 3, ]
    test.test_function(pathSum, inp, out)
    test.test_function(pathSum_opt, inp, out)
